from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import matplotlib.cm as cm

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal


def get_intrinsic_pc():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    """ load model
    """
    params = load_weights(rpl_config.model_pth)

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)
    
    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 1000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)
    rnn_state_old = rnn_state.copy()
    diff = jnp.abs(rnn_state - rnn_state_old)
    rnn_state_old = rnn_state.copy()
    diff_norm = jnp.linalg.norm(diff, axis=1)
    diff_norm_old = diff_norm.copy()
    norm_std = diff_norm.copy()

    rnn_state_init = rnn_state.copy()

    obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

    rnn_state_trajectory = []

    for t in range(rpl_config.life_duration):

        if t == rpl_config.probe_point:
            rnn_state_init = rnn_state.copy()

        progress_bar(t, rpl_config.life_duration)

        """ model forward 
        """
        rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
        diff = jnp.abs(rnn_state - rnn_state_old)
        rnn_state_old = rnn_state.copy()
        diff_norm = jnp.linalg.norm(diff, axis=1)
        norm_std = 0.4 * jnp.abs(diff_norm - diff_norm_old) + 0.6 * norm_std
        diff_norm_old = diff_norm.copy()

        rnn_state_trajectory.append(np.array(rnn_state).copy())
            
    print(rnn_state.shape)
    print(norm_std.shape)
    rnn_state_np = np.array(rnn_state)

    # 将 rnn_state_trajectory 展开成 rnn_state_np 的形状
    rnn_state_trajectory_np = np.array(rnn_state_trajectory)
    rnn_state_trajectory_np = rnn_state_trajectory_np.reshape(-1, rnn_state_trajectory_np.shape[-1])
    print("shape of rnn_state_trajectory_np: ", rnn_state_trajectory_np.shape)

    # 对 rnn_state_np 进行 PCA
    pca = PCA()
    # pca.fit(rnn_state_np)
    pca.fit(rnn_state_trajectory_np)

    # 打印 variance ratio
    print(pca.explained_variance_ratio_)

    rnn_state_np_pca = pca.transform(rnn_state_np)

    # 创建KMeans对象，指定聚类数为4
    kmeans = KMeans(n_clusters=4)
    # 对rnn_state_np_pca进行聚类
    kmeans.fit(rnn_state_np_pca)
    # 获取聚类中心的坐标
    cluster_centers = kmeans.cluster_centers_

    return pca, cluster_centers


def main():
    
    seq_len = 15
    
    n_samples = 1000

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--seq_len", type=int, default=15)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    seq_len = args.seq_len

    """ load model
    """
    params = load_weights(rpl_config.model_pth)
    
    nn_type = ''
    if rpl_config.nn_type == "vanilla":
        nn_type = "vanilla"
    elif rpl_config.nn_type == "gru":
        nn_type = "gru"

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)
    
    """ load task trajectories
    """
    slice_fn = "_" + nn_type + "_" + str(seq_len)

    # 载入 er12a 数据
    er12_data = np.load("./logs/tracelet_slices_and_chosen_er9_tracelet_and_obs_seq_low_ambiguity.npz")
    er12_tracelet_slices = er12_data["tracelet_slices"]
    er12_chosen_er9_tracelet = er12_data["chosen_er9_tracelet"]
    er12_obs_seq_low_ambiguity_final = er12_data["obs_seq_low_ambiguity_final"]

    print("shape of er12_tracelet_slices: ", er12_tracelet_slices.shape)
    print("shape of er12_chosen_er9_tracelet: ", er12_chosen_er9_tracelet.shape)
    print("shape of er12_obs_seq_low_ambiguity_final: ", er12_obs_seq_low_ambiguity_final.shape)
    
    # 载入 er9 数据
    er9_data = np.load("./logs/e_r9_data_gru.npz")
    er9_rnn_state_trajectory_set = er9_data["rnn_state_trajectory_set"]
    er9_tracelet_slices1 = er9_data["tracelet_slices1"]
    er9_tracelet_slices2 = er9_data["tracelet_slices2"]
    er9_tracelet_slices3 = er9_data["tracelet_slices3"]
    er9_tracelet_slices4 = er9_data["tracelet_slices4"]
    er9_tracelet_slices5 = er9_data["tracelet_slices5"]
    er9_seq_len = er9_data["seq_len"]
    er9_collect_capacity = er9_data["collect_capacity"]

    er9_tracelets = [er9_tracelet_slices1[0], er9_tracelet_slices2[0], er9_tracelet_slices3[0], er9_tracelet_slices4[0], er9_tracelet_slices5[0]]
    er9_tracelets = np.array(er9_tracelets)
    print("shape of er9_tracelets: ", er9_tracelets.shape)

    # 将 er9_rnn_state_trajectory_set 切分成5个部分，每个部分有 er9_collect_capacity 条轨迹
    er9_rnn_state_trajectory_seperate = []
    for i in range(5):
        er9_rnn_state_trajectory_seperate.append(er9_rnn_state_trajectory_set[i*er9_collect_capacity:(i+1)*er9_collect_capacity])
    er9_rnn_state_trajectory_seperate = np.array(er9_rnn_state_trajectory_seperate)
    print("shape of er9_rnn_state_trajectory_seperate: ", er9_rnn_state_trajectory_seperate.shape)

    er9_rnn_state_limit_ring_seperate = er9_rnn_state_trajectory_seperate[:,:,-2*seq_len:-1,:]
    print("shape of er9_rnn_state_limit_ring_seperate: ", er9_rnn_state_limit_ring_seperate.shape)

    # 计算 5 个类别的 cluster center
    er9_limit_ring_cluster_centers = []
    for i in range(5):
        cluster = er9_rnn_state_limit_ring_seperate[i]
        cluster = cluster.reshape(-1, cluster.shape[-1])
        cluster_center = np.mean(cluster, axis=0)
        er9_limit_ring_cluster_centers.append(cluster_center)
    er9_limit_ring_cluster_centers = np.array(er9_limit_ring_cluster_centers)
    print("shape of er9_limit_ring_cluster_centers: ", er9_limit_ring_cluster_centers.shape)

    """ 进行单周期、零初值扰动测试
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    n_samples = er12_obs_seq_low_ambiguity_final.shape[0]

    rnn_state = model.initial_state(n_samples)
    # k1 = npr.randint(0, 1000000)
    # rnn_state = model.initial_state_rnd(n_samples, k1)

    # 交换 er12_obs_seq_low_ambiguity_final 的第一维和第二维
    er12_obs_seq_low_ambiguity_final = np.swapaxes(er12_obs_seq_low_ambiguity_final, 0, 1)

    print("shape of er12_obs_seq_low_ambiguity_final: ", er12_obs_seq_low_ambiguity_final.shape)

    """ run forced dynamics
    """
    for kk in range(100):
        progress_bar(kk, 100)
        for i in range(er12_obs_seq_low_ambiguity_final.shape[0]):

            obs = er12_obs_seq_low_ambiguity_final[i]
            rnn_state, y1 = model_forward(params, rnn_state, obs, model)

    print("shape of rnn_state: ", rnn_state.shape)

    # 随机定义五种颜色，用于后续的 plt 绘制
    colors = ["red", "green", "blue", "orange", "purple"]

    # 测量 np.mean(rnn_state, axis=0) 距离 5 个 cluster center 哪个最近
    dist = []
    for i in range(5):
        dist.append(np.linalg.norm(np.mean(rnn_state, axis=0) - er9_limit_ring_cluster_centers[i]))
    dist = np.array(dist)
    min_dist_idx = np.argmin(dist)
    print("min_dist_idx: ", colors[min_dist_idx])

    """ visualization
    """
    er9_rnn_state_limit_ring_seperate_linear = er9_rnn_state_limit_ring_seperate.reshape(-1, er9_rnn_state_limit_ring_seperate.shape[-1])

    pca = PCA()
    pca.fit(er9_rnn_state_limit_ring_seperate_linear)

    rnn_state_center = np.mean(rnn_state, axis=0)
    rnn_state_center = rnn_state_center.reshape(1, -1) 

    rnn_state_pca = pca.transform(rnn_state_center)
    rnn_states_pca = pca.transform(rnn_state)

    # 将 er9_rnn_state_limit_ring_seperate 的所有轨迹在 PCA 空间中绘制出来
    fig = plt.figure(figsize=(12, 5))
    ax1 = fig.add_subplot(121, projection='3d')  # 3D plot
    ax2 = fig.add_subplot(122)  # 2D plot
    for i in range(5):
        cluster = er9_rnn_state_limit_ring_seperate[i]
        cluster = cluster.reshape(-1, cluster.shape[-1])
        cluster_pca = pca.transform(cluster)
        ax1.scatter(cluster_pca[:, 0], cluster_pca[:, 1], cluster_pca[:, 2], c=colors[i], alpha=0.1)

    # 使用大号圆点绘制 cluster center
    for i in range(5):
        cluster_center = er9_limit_ring_cluster_centers[i]
        cluster_center_pca = pca.transform(cluster_center.reshape(1, -1))
        ax1.scatter(cluster_center_pca[:, 0], cluster_center_pca[:, 1], cluster_center_pca[:, 2], c=colors[i], s=100)

    # 使用灰色圆点绘制 rnn_state_pca
    # ax1.scatter(rnn_state_pca[:, 0], rnn_state_pca[:, 1], rnn_state_pca[:, 2], c="grey", alpha=1, s=100)
    ax1.scatter(rnn_states_pca[:, 0], rnn_states_pca[:, 1], rnn_states_pca[:, 2], c="grey", alpha=0.1)

    # 绘制第二个图形
    ax2.plot(er9_tracelet_slices1[0, :, 0], er9_tracelet_slices1[0, :, 1], c=colors[0])
    ax2.plot(er9_tracelet_slices2[0, :, 0], er9_tracelet_slices2[0, :, 1], c=colors[1])
    ax2.plot(er9_tracelet_slices3[0, :, 0], er9_tracelet_slices3[0, :, 1], c=colors[2])
    ax2.plot(er9_tracelet_slices4[0, :, 0], er9_tracelet_slices4[0, :, 1], c=colors[3])
    ax2.plot(er9_tracelet_slices5[0, :, 0], er9_tracelet_slices5[0, :, 1], c=colors[4])

    # 在 ax2 上绘制 er12_tracelet_slices
    er12_tracelet_slices_reg = er12_tracelet_slices - er12_tracelet_slices[0]
    ax2.plot(er12_tracelet_slices_reg[:, 0], er12_tracelet_slices_reg[:, 1], c="grey")

    # 调整子图之间的间距
    plt.subplots_adjust(wspace=0.3)

    plt.show()


if __name__ == "__main__":
    main()